#include <condition_variable>
#include <functional>
#include <future>
#include <mutex>
#include <queue>
#include <thread>
#include <vector>
#include <stdexcept>
#include <memory>
#include <atomic>      // 原子操作支持
#include <type_traits> // 类型 traits 支持（std::decay、std::result_of）
#include <iostream>

class ThreadPool {
public:
	// 构造函数：初始化线程池并启动指定数量的线程
	// maxQueueSize为0表示无限制
	ThreadPool(size_t numThreads, size_t maxQueueSize = 0) :
	    stop(false),    // 停止标志（原子变量）
	    activeTasks(0), // 当前执行中的任务数（原子变量）
	    maxQueueSize(maxQueueSize) {

		if (numThreads == 0) {
			throw std::invalid_argument("Number of threads cannot be 0");
		}

		// 创建工作线程：每个线程运行workerThread函数
		threads.reserve(numThreads);
		for (size_t i = 0; i < numThreads; ++i) {
			threads.emplace_back(&ThreadPool::workerThread, this);
		}
	}

	// 禁止拷贝构造和赋值
	ThreadPool(const ThreadPool &)            = delete;
	ThreadPool &operator=(const ThreadPool &) = delete;

	// 任务入队（支持任意可调用对象及参数）
	// 返回future以便获取结果或异常（兼容C++11的std::result_of）
	template<class F, class... Args>
	auto enqueue(F &&f, Args &&...args) -> std::future<typename std::result_of<F(Args...)>::type> {

		// 推导任务返回类型（C++11兼容）
		using ReturnType = typename std::result_of<F(Args...)>::type;

		// 安全封装任务：用std::decay确保F是值类型，避免引用失效
		typedef typename std::decay<F>::type FuncType;
		auto                                 task = std::make_shared<std::packaged_task<ReturnType()>>(std::bind(static_cast<FuncType>(f), std::forward<Args>(args)...));

		// 获取任务的future（用于主线程等待结果）
		std::future<ReturnType> res = task->get_future();

		{
			std::unique_lock<std::mutex> lock(queue_mutex);

			// 线程池已停止，拒绝新任务
			if (stop) {
				throw std::runtime_error("Cannot add tasks to a stopped thread pool");
			}

			// 队列有大小限制时，等待队列有空位
			if (maxQueueSize > 0) {
				queue_not_full.wait(lock, [this] { return tasks.size() < maxQueueSize || stop; });
				// 等待后再次检查停止状态（避免停止后还加任务）
				if (stop) {
					throw std::runtime_error("Cannot add tasks to a stopped thread pool");
				}
			}

			// 将任务包装为无参Lambda，存入队列
			tasks.emplace([task]() { (*task)(); });
		}

		// 通知一个等待线程：有新任务了
		queue_not_empty.notify_one();
		return res;
	}

	// 查询当前正在执行的任务数量（原子操作，无需加锁）
	size_t activeCount() const { return activeTasks.load(); }

	// 查询等待中的任务数量（需要加锁保护队列）
	size_t pendingCount() const {
		std::unique_lock<std::mutex> lock(queue_mutex);
		return tasks.size();
	}

	// 阻塞等待所有任务完成（包括队列中的和正在执行的）
	void waitAll() {
		std::unique_lock<std::mutex> lock(queue_mutex);
		// 等待条件：队列空 + 无活动任务
		all_completed.wait(lock, [this] { return tasks.empty() && activeTasks.load() == 0; });
	}

	// 析构函数：安全关闭线程池（唤醒所有线程并等待退出）
	~ThreadPool() {
		// 1. 设置停止标志（锁内操作，确保同步）
		{
			std::unique_lock<std::mutex> lock(queue_mutex);
			stop = true;
		}

		// 2. 唤醒所有等待的线程（队列非空 + 队列未满的条件变量都要通知）
		queue_not_empty.notify_all();
		queue_not_full.notify_all();

		// 3. 等待所有工作线程退出
		for (std::thread &worker: threads) {
			if (worker.joinable()) {
				worker.join();
			}
		}
	}

private:
	// 工作线程的核心逻辑
	void workerThread() {
		while (true) {
			std::function<void()> task;
			{
				std::unique_lock<std::mutex> lock(queue_mutex);

				// 等待条件：队列非空 或 线程池停止（避免虚假唤醒）
				queue_not_empty.wait(lock, [this] { return stop || !tasks.empty(); });

				// 线程池已停止且队列空，线程退出
				if (stop && tasks.empty()) {
					return;
				}

				// 取出任务（移动语义，减少拷贝）
				task = std::move(tasks.front());
				tasks.pop();
			}

			// --- 任务执行阶段（无锁） ---
			++activeTasks; // 原子递增：当前任务开始执行

			try {
				task(); // 执行任务
			} catch (const std::exception &e) {
				// 捕获标准异常，打印错误信息（避免静默失败）
				std::cerr << "[Thread Pool] Task execution exception: " << e.what() << std::endl;
			} catch (...) {
				// 捕获未知异常，打印提示
				std::cerr << "[Thread Pool] Task execution unknown exception" << std::endl;
			}

			--activeTasks; // 原子递减：当前任务执行完成
			// --- 检查是否所有任务完成 ---
			{
				std::unique_lock<std::mutex> lock(queue_mutex);
				// 条件：队列空 + 无活动任务 → 通知waitAll
				if (tasks.empty() && activeTasks.load() == 0) {
					all_completed.notify_all();
				}
			}
		}
	}

	// 成员变量说明
	std::atomic<bool>                 stop;         // 线程池停止标志（原子操作）
	std::atomic<size_t>               activeTasks;  // 正在执行的任务数（原子操作）
	size_t                            maxQueueSize; // 任务队列最大容量
	std::vector<std::thread>          threads;      // 工作线程集合
	std::queue<std::function<void()>> tasks;        // 待执行任务队列（无参Lambda包装）

	mutable std::mutex      queue_mutex;     // 任务队列的互斥锁
	std::condition_variable queue_not_empty; // 队列非空条件变量（唤醒工作线程）
	std::condition_variable queue_not_full;  // 队列未满条件变量（唤醒入队线程）
	std::condition_variable all_completed;   // 所有任务完成条件变量（唤醒waitAll）
};

// ===================== 测试用例（保持原逻辑） =====================
#include <iostream>
#include <iomanip>
#include <chrono>
#include <random>
#include <string>

// 简单测试：验证基本功能
void test_simple() {
	std::cout << "=== 简单功能测试 === " << std::endl;
	ThreadPool pool(4);

	// 入队普通函数
	pool.enqueue([](const std::string &name) { std::cout << "Hello " << name << std::endl; }, "World");

	pool.enqueue([](int a, int b) { std::cout << a << " + " << b << " = " << a + b << std::endl; }, 2, 3);

	// 入队lambda并获取结果
	auto future = pool.enqueue([]() -> int {
		std::this_thread::sleep_for(std::chrono::milliseconds(100));
		return 42;
	});

	std::cout << "Lambda返回值: " << future.get() << std::endl;
	pool.waitAll();
}

// 并发计数测试：验证线程安全和计数准确性
void test_concurrent_counting() {
	std::cout << "=== 并发计数测试 === " << std::endl;
	size_t     n = std::thread::hardware_concurrency();
	ThreadPool pool(n > 0 ? n : 4); // 自动适配硬件并发数

	auto      counter    = std::make_shared<size_t>(0);
	auto      mtx        = std::make_shared<std::mutex>(); // 保护计数器的互斥锁
	const int task_count = 50;

	for (int i = 0; i < task_count; ++i) {
		pool.enqueue([counter, mtx, &pool, i]() {
			std::this_thread::sleep_for(std::chrono::microseconds(1000 + rand() % 1000));

			std::lock_guard<std::mutex> lock(*mtx);
			++(*counter);

			std::cout << "线程ID: " << std::this_thread::get_id() << " 任务: " << std::setw(2) << std::setfill('0') << i << " 活动任务: " << pool.activeCount()
			          << " 等待任务: " << pool.pendingCount() << " 计数: " << *counter << std::endl;
		});
	}

	pool.waitAll();
	std::cout << "最终计数: " << *counter << " (预期: " << task_count << ")" << std::endl;
}

// 测试返回值：验证future的正确性
void test_future_results() {
	std::cout << "=== 测试返回值 === " << std::endl;
	ThreadPool pool(2);

	std::vector<std::future<int>> futures;
	for (int i = 0; i < 5; ++i) {
		futures.emplace_back(pool.enqueue([i]() -> int {
			std::this_thread::sleep_for(std::chrono::milliseconds(100 * i));
			return i * i;
		}));
	}

	for (auto &fut: futures) {
		std::cout << "计算结果: " << fut.get() << std::endl;
	}
}

// 测试队列限制：验证队列满时的阻塞行为
void test_queue_limit() {
	std::cout << "=== 测试任务队列限制 === " << std::endl;
	ThreadPool pool(1, 5); // 1个线程，队列最大5个任务

	auto start = std::chrono::high_resolution_clock::now();

	// 添加10个任务（队列满时会阻塞）
	for (int i = 0; i < 10; ++i) {
		pool.enqueue([i]() {
			std::this_thread::sleep_for(std::chrono::milliseconds(100));
			std::cout << "完成任务 " << i << std::endl;
		});

		std::cout << "添加任务 " << i << "，当前等待任务数: " << pool.pendingCount() << std::endl;
	}

	pool.waitAll();
	auto                          end     = std::chrono::high_resolution_clock::now();
	std::chrono::duration<double> elapsed = end - start;
	std::cout << "总耗时: " << elapsed.count() << "秒 (预期约1秒)" << std::endl;
}

// 编译： g++ -std=c++11 -pthread thread_pool_example.cpp -o thread_pool_example.exe
int main() {
	srand(time(nullptr)); // 初始化随机数种子

	test_simple();
	test_concurrent_counting();
	test_future_results();
	test_queue_limit();

	return 0;
}
